import requests
import sqlite3
import pandas as pd
from datetime import datetime
import time
import os

# Cookie 管理功能
def get_saved_cookie():
    """尝试从 cookie.txt 文件加载 Cookie"""
    cookie_file = "cookie.txt"
    if os.path.exists(cookie_file):
        try:
            with open(cookie_file, "r") as f:
                return f.read().strip()
        except:
            return None
    return None

def save_new_cookie(cookie):
    """保存新 Cookie 到文件"""
    try:
        with open("cookie.txt", "w") as f:
            f.write(cookie)
        return True
    except:
        return False

# 用户输入 cookie (优先使用保存的 cookie)
saved_cookie = get_saved_cookie()
if saved_cookie:
    cookie_input = saved_cookie
    print("使用保存的 Cookie")
else:
    cookie_input = input("首次使用，请输入 Cookie 值：")
    save_new_cookie(cookie_input)
    
db_folder = os.environ.get('TLL_database')
db_path = os.path.join(db_folder, 'user_data.db')
conn = sqlite3.connect(db_path)

headers = {
    "user-agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/136.0.0.0 Safari/537.36 Edg/136.0.0.0",
    "cookie": cookie_input
}

def fetch_page(page_num):
    url = f'https://account.italent.cn/Contacts/staffV2?page_num={page_num}&page_size=50&department_id=0&_v=1747805471108&_qsrcapp=AccountSite&quark_s=d6edd30be119e0290e1ed63723a14492ba9e201588f585ca50244637e34672e5'

    try:
        resp = requests.get(url, headers=headers, timeout=10)
        resp.raise_for_status()
        
        # 检查 Cookie 是否失效
        if resp.status_code == 401 or resp.status_code == 403:
            print(f"Page {page_num}: Cookie 已失效")
            return None, True  # 返回需要更新 Cookie 标志
        
        try:
            resp_json = resp.json()
        except ValueError:
            print(f"Page {page_num}: 响应内容不是合法的 JSON")
            return None, False

        if not isinstance(resp_json, dict):
            print(f"Page {page_num}: 返回数据不是字典类型，无法解析")
            return None, False

        if 'Data' in resp_json and isinstance(resp_json['Data'], dict) and 'users' in resp_json['Data']:
            return resp_json['Data']['users'], False
        else:
            print(f"Page {page_num}: 数据中缺少 'Data' 或 'users'")
            return None, False

    except requests.exceptions.RequestException as e:
        print(f"Page {page_num} 请求失败: {e}")
        return None, False


def save_to_db_with_pandas(df, table_name='users'):
    df.to_sql(table_name, conn, if_exists='append', index=False)
    print(f"已将当前页数据通过 Pandas 写入 SQLite 表 '{table_name}'")


page = 1
all_data = []
exit_flag = False  # 退出标志

while not exit_flag:
    print(f"Fetching page {page}...")
    userjson, need_new_cookie = fetch_page(page)

    # 处理 Cookie 失效情况
    if need_new_cookie:
        print("Cookie 已失效，请输入新的 Cookie 值")
        cookie_input = input("新的 Cookie 值：")
        save_new_cookie(cookie_input)  # 保存新 Cookie
        headers['cookie'] = cookie_input  # 更新当前请求中的 Cookie
        
        # 使用新 Cookie 重试当前页面
        print("使用新 Cookie 重试当前页...")
        userjson, need_new_cookie = fetch_page(page)
        if need_new_cookie:
            print("新 Cookie 仍然无效，程序终止")
            exit_flag = True
            continue

    # 保持原始分页处理逻辑不变
    if not userjson:  # 如果获取不到数据，跳出循环
        print(f"Page {page}: 未获取到数据，退出循环")
        exit_flag = True
        continue

    current_time = datetime.now().strftime('%Y-%m-%d %H:%M:%S')  # 获取当前时间
    records = []
    for user in userjson:
        info_dict = {item['key']: item['value'] for item in user.get('infoValues', [])}
        record = {
            'userId': user['userId'],
            'name': user['name'],
            'gender': info_dict.get('Sex'),
            'department': user.get('departmentPath', ''),
            'email': info_dict.get('Email'),
            'phone': info_dict.get('Mobile'),
            'import_date': current_time,
            'remark': ""  # 默认为空备注
        }
        records.append(record)

    df = pd.DataFrame(records)
    save_to_db_with_pandas(df)
    all_data.append(df)

    page += 1
    time.sleep(1)  # 避免请求过快

# 关闭数据库连接
conn.close()
print("程序执行完毕")